import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from io import BytesIO

# 设置应用标题和描述
st.title('COVID-19 数据分析应用')
st.write('上传CSV格式的疫情数据，进行数据分析和可视化。')

# 数据上传
uploaded_file = st.file_uploader("请选择一个CSV文件上传", type="csv")
if uploaded_file is not None:
    # 数据预处理
    try:
        data = pd.read_csv(uploaded_file)
        # 确保数据按日期排序
        data['dateId'] = pd.to_datetime(data['dateId'])
        data.sort_values('dateId', inplace=True)
        data.dropna(inplace=True)  # 去除空值

        # 数据分析
        # 总感染趋势
        total_infection_trend = data.groupby('dateId')['confirmedCount'].sum().reset_index()

        # 各城市感染曲线
        city_infection_trend = data.groupby(['dateId', 'cityName'])['confirmedCount'].sum().reset_index()

        # 风险区域识别
        threshold = st.slider('设置感染人数阈值', min_value=0, max_value=int(data['confirmedCount'].max()), value=50)
        high_risk_areas = city_infection_trend[city_infection_trend['confirmedCount'] > threshold]

        # 图表展示
        fig, ax = plt.subplots()
        total_infection_trend.plot(x='dateId', y='confirmedCount', ax=ax, label='总感染趋势')
        ax.set_title('新冠病毒总感染趋势')
        ax.set_xlabel('日期')
        ax.set_ylabel('累计确诊人数')
        st.pyplot(fig)

        fig, ax = plt.subplots(figsize=(15, 5))
        for city, group in city_infection_trend.groupby('cityName'):
            if city in high_risk_areas['cityName'].unique():
                ax.plot(group['dateId'], group['confirmedCount'], label=city, linestyle='--')
            else:
                ax.plot(group['dateId'], group['confirmedCount'], label=city)
        ax.set_title('各城市感染曲线')
        ax.set_xlabel('日期')
        ax.set_ylabel('累计确诊人数')
        ax.legend()
        st.pyplot(fig)

        # 用户交互
        # 时间范围选择
        start_date = st.date_input("选择开始日期", value=data['dateId'].min().to_pydatetime().date())
        end_date = st.date_input("选择结束日期", value=data['dateId'].max().to_pydatetime().date())

        # 将选择的日期转换为 pandas datetime64 类型
        start_date = pd.to_datetime(start_date)
        end_date = pd.to_datetime(end_date)

        # 使用 pandas datetime64 类型进行过滤
        filtered_data = data[(data['dateId'] >= start_date) & (data['dateId'] <= end_date)]
        # 结果输出
        # 下载按钮
        if st.button('下载分析结果图表'):
            buf = BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            st.download_button(label="Download image", data=buf, file_name="infection_trend.png", mime="image/png")
    except Exception as e:
        st.error(f"发生错误：{e}")
